import argparse
from pathlib import Path
import re
from collections import defaultdict
import shutil

parser = argparse.ArgumentParser(description = "card注释")
parser.add_argument("-i", dest = "input_file", type = str, required = True)
parser.add_argument("-outdir", dest = "outdir", type = str, required = True)
parser.add_argument("-s", dest = "sample_id", type = str, required = True)
parser.add_argument("-c", dest = "card_db", type = str, nargs = '?')
parser.add_argument("-a", dest = "card_aro", type = str, nargs = '?')
args = parser.parse_args()
       
input_file = args.input_file
outdir = Path(args.outdir)
sample_id = args.sample_id
pydir = Path(__file__).parent
# 抗性机制中文
antibiotic_resistant = pydir / 'antibiotic-resistant.txt'
antibiotic_resistant_dict = defaultdict(str)
with open(antibiotic_resistant, 'r') as ares:
    for line in ares:
        line = line.strip().split('\t')
        antibiotic_resistant_dict[line[0]] = line[1]
# class 中文
antibiotic = pydir / 'antibiotic.txt'
antibiotic_dict = defaultdict(str)
with open(antibiotic, 'r') as anti:
    for line in anti:
        line = line.strip().split('\t')
        antibiotic_dict[line[0]] = line[1]
# tax中文
microbio = pydir / 'microbio-name.txt'
microbio_dict = defaultdict(str)
with open(microbio, 'r') as micr:
    for line in micr:
        line = line.strip().split('\t')
        if len(line) == 2:
            microbio_dict[line[0]] = line[1]
# ARO总长度/基因家族
if args.card_aro:
    card_aro = Path(args.card_aro)
else:
    card_aro = pydir / '../ref/aro_len.xls'
card_aro_len = defaultdict(int)
card_gene_family = defaultdict(str)
with open(card_aro, 'r') as caro:
    for line in caro:
        line = line.strip().split('\t')
        card_aro_len[line[0]] = line[1]
        card_gene_family[line[0]] = line[-1]

# card注释文件
if args.card_db:
    card = Path(args.card_db)
else:
    card = pydir / 'aro_index_tax.tsv'
card_anno = defaultdict(lambda: defaultdict(str))
with open(card, 'r') as card_db:
    line1 = card_db.readline().strip().split('\t')
    for line in card_db:
        line = line.strip().split('\t')
        for key,value in enumerate(line):
            card_anno[line[0]][line1[key]] = value

# 记录目标序列最优比对的结果
card_num2 = defaultdict(lambda: defaultdict(int))

# CIGAR
# 判断序列剪切是否因为参考序列头尾有关
m = re.compile(r'(\d+)(\D+)')
clipping = ['S', 'H']
def check(line, start, end):
    cigar = m.findall(line)
    clip = 0
    pos = 1
    for cell in cigar:
        if cell[1] in clipping:
            # 如果序列头被剪切，允许1碱基错配前提下判断
            if pos == 1:
                if int(start) <= 2:
                    clip = int(clip) + int(start) - 1
                else:
                    # 起始位点在参考序列中间处
                    clip = int(clip) + int(cell[0])
            else:
                if int(start) < int(end):
                    if pos < len(cigar):
                        clip = int(clip) + int(cell[0])
                        start = int(start) + int(cell[0])
                    else:
                        # 如果在参考序列尾部
                        if int(start) + 1 < int(end):
                            clip = int(clip) + int(cell[0])
                            start = int(start) + int(cell[0])
                        else:
                            # 这种情况只会出现在起始位置为参考序列最后1bp上，错配只需要为1
                            if int(start) == int(end):
                                clip = 1 + int(clip)
            if int(clip) > 1:
                return(-1)
        else:
            if cell[1] != 'M':
                # 如果是插入，起始位点不增加，只增加clip数
                if cell[1] == 'I':
                    clip = int(clip) + int(cell[0])
                else:
                    # 如果不为I，检查是否因为参考序列尾部造成的
                    if int(start) + 1 < int(end):
                        clip = int(clip) + int(cell[0])
                        start = int(start) + int(cell[0])
                if int(clip) > 1:
                    return(-1)
            else:
                start = int(start) + int(cell[0])
        pos += 1
    return(int(clip))

# 记录ARO覆盖度
aro = defaultdict(lambda: defaultdict(str))
aro_num = defaultdict(int)
clipping = ['S', 'H']
def coverage(line, start, id):
    cigar = m.findall(line)
    pos = 1
    for cell in cigar:
        if cell[1] in clipping:
            #如果不是在头部进行的剪切，将起始位置修改
            if pos != 1:
                start = int(start) + int(cell[0])
        else:
            if cell[1] != 'M':
                if cell[1] != 'I':
                    start = int(start) + int(cell[0])
            else:
                for i in range(int(start), int(start) + int(cell[0])):                    
                    if int(i) not in aro[id]:
                        aro[id][int(i)] = 1
                        aro_num[id] += 1
                start = int(start) + int(cell[0])
        pos += 1

with open(input_file, 'r') as ifile:
    for line in ifile:
        line = line.strip().split('\t')
        if line[5] == '*':
            continue
        # 记录覆盖位置
        #raw_name = re.search(r'.*?(ARO:.*)[_\|]+.*?' , line[2])
        #aro_id = raw_name.group(1)
        aro_id = line[2]
        decide = check(line[5], line[3], card_aro_len[line[2]])
        if int(decide) < 0:
            continue
        else:
            nm = re.search(r'NM:i:(.*)', line[11])
            nm = nm.group(1)
            nm = int(decide) + int(nm)
            if int(nm) <= 1:
                # 处理覆盖度
                coverage(line[5], line[3], aro_id)
                card_num2[aro_id]['best'] += 1
                xa = re.search(r'XA:Z:(.*)' , line[-1])
                if xa:
                    xa = xa.group(1).split(';')
                    for key, value in enumerate(xa):
                        if key == (len(xa) - 1):
                            continue
                        xa1 = value.split(',')
                        xa1[1] = xa1[1].replace('+','').replace('-','')
                        decide = check(xa1[2], xa1[1], card_aro_len[xa1[0]])
                        if int(decide) >= 0:
                            nm = int(decide) + int(xa1[-1])
                            if nm <= 1:
                                coverage(xa1[1], xa1[2], xa1[0])
                                card_num2[xa1[0]]['best'] += 1



output_file1 = outdir / f'{sample_id}.gene_mapping_data_bak.txt'
output_file2 = outdir / f'{sample_id}.gene_mapping_data_bak.xls'
with open(output_file1 , 'w') as ofile1 , open(output_file2 , 'w') as ofile2:
    ofile2.write('ID\tgene\tcount\tcoverage\tdrug_cn\tmed_cn\tspecies_cn\tdrug_en\tmed_en\tspecies_en\n')
    ofile1.write('ID\tgene\tcount\tcoverage\tdrug\tmed\tspecies\n')
    for key in card_num2.keys():
        # class转换成中文        
        class1 = card_anno[key]["Drug Class"].split(';')
        class_chi = []        
        for key22, value22 in enumerate(class1):
            if value22 in antibiotic_dict:
                class_chi.append(antibiotic_dict[value22])
            else:
                class_chi.append(value22)       
        # 耐药机制转换成中文
        mechanism = card_anno[key]["Resistance Mechanism"].split(';')
        mechanism_chi  = []
        for key22, value22 in enumerate(mechanism):
            if value22 in antibiotic_resistant_dict:
                mechanism_chi.append(antibiotic_resistant_dict[value22])
            else:
                mechanism_chi.append(value22)
        # 物种转换成中文
        microbio = card_anno[key]["NCBI Taxonomy Name"].split(';')
        microbio_chi = []
        for key22, value22 in enumerate(microbio):
            value33= value22.split(' ')
            if len(value33) > 2:
                value22 = f'{value33[0]} {value33[1]}'
            if value22 in microbio_dict:
                microbio_chi.append(microbio_dict[value22])
            else:
                microbio_chi.append(value22)

        ratio = str('%.2f' % ( int(aro_num[key]) * 100 / int(card_aro_len[key]) ))
        # 如果class为单纯的disinfecting agents and intercalating dyes，则删除此ARO
        if card_anno[key]["Drug Class"] == 'disinfecting agents and intercalating dyes':
            continue
        if card_anno[key]["NCBI Taxonomy Name"].find(r'Plasmid') > -1:
            continue
        numb2 = f'\
{sample_id}\t{card_anno[key]["ARO Name"]}\t{card_num2[key]["best"]}\t{ratio}\t\
{";".join(class_chi)}\t{";".join(mechanism_chi)}\t{";".join(microbio_chi)}\t\
{card_anno[key]["Drug Class"]}\t{card_anno[key]["Resistance Mechanism"]}\t{card_anno[key]["NCBI Taxonomy Name"]}\
'
        ofile2.write(f'{numb2}\n')
        numb1 = f'\
{sample_id}\t{card_anno[key]["ARO Name"]}\t{card_num2[key]["best"]}\t{ratio}\t\
{card_anno[key]["Drug Class"]}\t{card_anno[key]["Resistance Mechanism"]}\t{card_anno[key]["NCBI Taxonomy Name"]}\
'
        ofile1.write(f'{numb1}\n')

        


